import tensorflow as tf
from bert import modeling
import os
import create_input
import tokenization
import numpy as np
from config import  *
from rebuild_bert import *

prior_num = 0
if __name__=='__main__':

    x_train = np.load('.\\restore\\x_train.npy',allow_pickle=True)
    x_test = np.load('.\\restore\\x_test.npy',allow_pickle=True)
    y_train = np.asarray(np.load('.\\restore\\y_train.npy'),np.int32)
    y_test = np.asarray(np.load('.\\restore\\y_test.npy'),np.int32)
    #重建网络输出 在原有的bert模型上更改全连接层输出实现
    model = rebuild_bert(bert_config=modeling.BertConfig.from_json_file(FLAGS.bert_config_path),batch_size=FLAGS.batch_size,
                         max_seq_length=FLAGS.max_seq_length, is_training=True,categories=FLAGS.categories,learning_rate=0.00005)
    #加载要预测的数据
    model.load_data(x_train=x_train,x_test=x_test,vocab_file=FLAGS.vocab_path)
    #加载自己fine tuning的模型
    init_checkpoint = "D:\\data_mining_bert\saved_models\model_iter160ac=0.7182017543859649.ckpt"  # 自己保存的模型
    init_checkpoint = FLAGS.init_checkpoint_path#第一次训练用这个 FLAGS.init_checkpoint_path
    model.train(x_train=x_train,x_test=x_test,y_train=y_train,y_test=y_test,
                init_checkpoint=init_checkpoint,iter_num=1000,iter_per_valid=100,prior_iter_num=0)


